#!/usr/bin/env python3.9
# coding=utf-8

"""
Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
Description: python api test.
Author: MindX SDK
Create: 2022
History: NA
"""

import logging
import time
from multiprocessing import JoinableQueue, Queue
import traceback
import numpy as np
import cv2
from mindx.sdk import base
from mindx.sdk.base import Image, ImageProcessor, Size, Tensor

from utils import car_post_process
from config import *
import logging.config
logging.config.fileConfig("logger.conf")
logger = logging.getLogger("pipeline")

# 颜色识别
class CarUnionColorRec(object):
    def __init__(self, rec_model_path, batch_size, device_id, pid):
        self.rec_model = base.model(rec_model_path, deviceId=device_id)
        self.batch_size = batch_size
        self.device_id = device_id
        self.pid = pid

        self.INFER_WAIT_TIME = 10

    def rec_preprocess(self, img):
        height = 224
        width = 224
        # resize
        img = cv2.resize(img, (width, height))
        img = img[:, :, ::-1]  # BGR转RGB, 车牌颜色模型通道数在最后不用transpose
        img = img.astype(np.float32)

        # 归一化
        img /= 127.5
        img -= 1.
        return img

    def infer(self, input: Queue, output: Queue):
        try:
            ImageProcessor(self.device_id)
            count = 0
            img_arr_list = []
            input_data_list = []

            while True:
                input_data = input.get()
                input_data_list.append(input_data)
                img_input = input_data["det"][0]
                if img_input is None:
                    continue

                img = img_input
                count += 1

                img = self.rec_preprocess(img)
                img_arr_list.append(img)

                if len(img_arr_list) < self.batch_size:
                    continue

                tf_imgs = np.array(img_arr_list)
                img_tensor = Tensor(tf_imgs)
                img_tensor.to_device(self.device_id)
                img_tensor_vec = [img_tensor]

                # 4. rec model infer
                output_tensors = self.rec_model.infer(img_tensor_vec)[0]
                output_tensors.to_host()
                output_tensors = np.array(output_tensors)
                # 5. rec model postprocess
                for i, output_tensor in enumerate(output_tensors):
                    # 传入下一节点
                    post_data = input_data_list[i]
                    top_score = np.max(output_tensors[i], axis=0)
                    top_index = np.argmax(np.array(output_tensors[i]), axis=0)
                    if top_score > 0.4:
                        post_data["color"] = MODEL_CAR_COLOR_LABEL[top_index]
                    else:
                        post_data["color"] = "-"
                    output.put(post_data, block=True)

                ### process result
                img_arr_list = []
                input_data_list = []
        except Exception as e:
            logger.error(f"{self.pid} pid car color process2_pattern_rec failed")
            traceback.print_exc()


# 车辆类别
class CarUnionPropRec(object):
    def __init__(self, rec_model_path, batch_size, device_id, pid, threshold=0.5):
        self.rec_model = base.model(rec_model_path, deviceId=device_id)
        self.batch_size = batch_size
        self.device_id = device_id
        self.pid = pid
        self.INFER_WAIT_TIME = 10
        self.threshold = threshold

    def rec_preprocess(self, img):
        height = 224
        width = 224

        # resize
        imh, imw = img.shape[:2]
        m = min(imh, imw)  # min dimension
        top, left = (imh - m) // 2, (imw - m) // 2
        img = cv2.resize(img[top:top + m, left:left + m], (width, height), interpolation=cv2.INTER_LINEAR)
        # img = cv2.resize(img, (width, height))

        img = img.transpose((2, 0, 1))  # BGR 转RGB
        img = img.astype(np.float32)

        # 归一化
        img /= 255.0

        # 标准化
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        mean = np.array(mean).reshape(3, 1, 1)
        std = np.array(std).reshape(3, 1, 1)
        img = (img-mean)/std
        img = img.astype(np.float32)
        img = np.ascontiguousarray(img)

        return img

    def infer(self, input: Queue, output: Queue):
        def sigmoid(z):
            return 1/(1 + np.exp(-z))
        try:
            ImageProcessor(self.device_id)
            count = 0
            img_arr_list = []
            input_data_list = []

            while True:
                input_data = input.get()
                input_data_list.append(input_data)
                img_input = input_data["det"][0]
                if img_input is None:
                    continue

                img = img_input
                count += 1
                img = self.rec_preprocess(img)
                img_arr_list.append(img)
                if len(img_arr_list) < self.batch_size:
                    continue

                tf_imgs = np.array(img_arr_list)
                img_tensor = Tensor(tf_imgs)
                img_tensor.to_device(self.device_id)
                img_tensor_vec = [img_tensor]

                # 4. rec model infer
                output_tensors = self.rec_model.infer(img_tensor_vec)[0]
                output_tensors.to_host()
                output_tensors = np.array(output_tensors)

                # 5. rec model postprocess
                for i, output_tensor in enumerate(output_tensors):
                    # 传入下一节点
                    post_data = input_data_list[i]
                    top_score = np.max(output_tensors[i], axis=0)
                    top_index = np.argmax(np.array(output_tensors[i]), axis=0)
                    if top_score> 0.5:
                        post_data["prop"] = MODEL_CAR_CLASS_LABEL[top_index]
                    else:
                        post_data["prop"] = "-"
                    output.put(post_data, block=True)

                ### process result
                img_arr_list = []
                input_data_list = []
        except Exception as e:
            logger.error(f"{self.pid} pid car prop process2_pattern_rec failed")
            traceback.print_exc()


# 车辆品牌
class CarUnionBrandRec(object):
    def __init__(self, rec_model_path, batch_size, device_id, pid, threshold=0.5):
        self.rec_model = base.model(rec_model_path, deviceId=device_id)
        self.batch_size = batch_size
        self.device_id = device_id
        self.pid = pid
        self.INFER_WAIT_TIME = 10
        self.threshold = threshold
        
    def rec_preprocess(self, img):
        height = 224
        width = 224

        # resize
        imh, imw = img.shape[:2]
        m = min(imh, imw)  # min dimension
        top, left = (imh - m) // 2, (imw - m) // 2
        img = cv2.resize(img[top:top + m, left:left + m], (width, height), interpolation=cv2.INTER_LINEAR)

        img = img.transpose((2, 0, 1))  # BGR
        img = img.astype(np.float32)

        # 归一化
        img /= 255.0

        # 标准化
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        mean = np.array(mean).reshape(3, 1, 1)
        std = np.array(std).reshape(3, 1, 1)
        img = (img-mean)/std
        img = img.astype(np.float32)
        img = np.ascontiguousarray(img)

        return img

    def infer(self, input: Queue, output: Queue):
        def sigmoid(z):
            return 1/(1 + np.exp(-z))
        try:
            ImageProcessor(self.device_id)
            count = 0
            img_arr_list = []
            input_data_list = []

            while True:
                input_data = input.get()
                input_data_list.append(input_data)
                img_input = input_data["det"][0]
                if img_input is None:
                    continue

                img = img_input
                count += 1
                img = self.rec_preprocess(img)
                img_arr_list.append(np.array(img))

                if len(img_arr_list) < self.batch_size:
                    continue

                img_tensor = Tensor(np.array(img_arr_list))

                img_tensor.to_device(self.device_id)
                img_tensor_vec = [img_tensor]

                # 4. rec model infer
                output_tensors = self.rec_model.infer(img_tensor_vec)[0]
                output_tensors.to_host()
                output_tensors = np.array(output_tensors)

                # 5. rec model postprocess
                for i, output_tensor in enumerate(output_tensors):
                    # 传入下一节点
                    post_data = input_data_list[i]
                    top_score = np.max(output_tensors[i], axis=0)
                    top_index = np.argmax(np.array(output_tensors[i]), axis=0)
                    if top_score > 0.4:
                        post_data["brand"] = MODEL_CAR_BRAND_LABEL[top_index]
                    else:
                        post_data["brand"] = "-"
                    output.put(post_data, block=True)

                ### process result
                img_arr_list = []
                input_data_list = []
        except Exception as e:
            logger.error(f"{self.pid} pid car brand process2_pattern_rec failed")
            traceback.print_exc()


def car_union_color_recognize(rec_model_path, batch_size, device_id, pid,
                              input: Queue,
                              output: Queue,
                              model_loading: list):

    model = CarUnionColorRec(rec_model_path, batch_size, device_id, pid)
    model_loading.append(True)
    model.infer(input, output)


def car_union_prop_recognize(rec_model_path, batch_size, device_id, pid,
                             input: Queue,
                             output: Queue,
                             model_loading: list):
    model = CarUnionPropRec(rec_model_path, batch_size, device_id, pid)
    model_loading.append(True)
    model.infer(input, output)


def car_union_brand_recognize(rec_model_path, batch_size, device_id, pid,
                              input: Queue,
                              output: Queue,
                              model_loading: list):
    model = CarUnionBrandRec(rec_model_path, batch_size, device_id, pid)
    model_loading.append(True)
    model.infer(input, output)